{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tracer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-11-24T04:20:48.699252Z",
     "start_time": "2022-11-24T04:20:48.694265Z"
    }
   },
   "outputs": [],
   "source": [
    "import cv2 as cv\n",
    "import numpy as np\n",
    "import imutils,math\n",
    "from datetime import datetime,timedelta\n",
    "import matplotlib as mpl\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-11-24T04:20:48.714212Z",
     "start_time": "2022-11-24T04:20:48.701247Z"
    }
   },
   "outputs": [],
   "source": [
    "global V_FF, DURATION, cales, scale, transf_para, FF, BIN_THR, V_MIN, V_MAX, HEAT_SIZE, ROIS\n",
    "V_FF = 30  # Frame frequency\n",
    "DURATION = V_FF * 5*60  # Frame frequency * seconds * minutes, default\n",
    "FY = 1  # resize factor\n",
    "FX = 1  # resize factor\n",
    "scale = 0\n",
    "HEIGHT = 0\n",
    "WIDTH = 0\n",
    "ROIS = []\n",
    "FF = 2  # Frame frequency\n",
    "BIN_THR = 50  # Binary threshold\n",
    "V_MIN = 1   # Set minimum value of velocity\n",
    "V_MAX = 90  # Set maxium value of velocity\n",
    "HEAT_SIZE = 16   # Set the size of heatmap"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-11-24T04:20:48.744132Z",
     "start_time": "2022-11-24T04:20:48.716206Z"
    }
   },
   "outputs": [],
   "source": [
    "def on_EVENT_LBUTTONDOWN(event, x, y, flags, param):\n",
    "    global ROIS\n",
    "    if event == cv.EVENT_LBUTTONDOWN:\n",
    "        if len(ROIS) < 2:\n",
    "            ROIS.append((x, y))\n",
    "        else:\n",
    "            ROIS = []\n",
    "\n",
    "def get_distance(p0, p1):\n",
    "    s = math.pow((p0[0] - p1[0]), 2) + math.pow((p0[1] - p1[1]), 2)\n",
    "    s = math.sqrt(s)\n",
    "    return s\n",
    "\n",
    "def get_roi(capture, video, no, f):\n",
    "    rot = 0\n",
    "    capture.set(cv.CAP_PROP_POS_FRAMES, f)\n",
    "    preret, prev_img = capture.read()\n",
    "    prev_img = cv.cvtColor(prev_img, cv.COLOR_RGB2GRAY)\n",
    "    (HEIGHT, WIDTH) = prev_img.shape\n",
    "    prev_img = cv.resize(prev_img, (0, 0), fx=FX, fy=FY)\n",
    "    (HEIGHT, WIDTH) = prev_img.shape\n",
    "    rot_center = (WIDTH // 2, HEIGHT // 2)\n",
    "    cv.namedWindow(\"Preview\")\n",
    "    cv.imshow(\"Preview\", prev_img)\n",
    "\n",
    "    cv.imshow(\"SetROI\", prev_img)\n",
    "    cv.setMouseCallback(\"SetROI\", on_EVENT_LBUTTONDOWN)\n",
    "\n",
    "    while 1:\n",
    "        if cv.waitKey(0) == 13:\n",
    "            break\n",
    "    cv.line(prev_img, (np.int(ROIS[0][0]), np.int(ROIS[0][1])),\n",
    "            (np.int(ROIS[1][0]), np.int(ROIS[1][1])), (255, 0, 0), 3)\n",
    "    #     Draw line\n",
    "    cv.destroyAllWindows()\n",
    "\n",
    "    roi_h = np.min([WIDTH, HEIGHT])\n",
    "    x, y, a = [0, 0, roi_h]\n",
    "\n",
    "    # select roi and rotation\n",
    "    while 1:\n",
    "        roi_img = np.zeros((HEIGHT, WIDTH), dtype=np.uint8)\n",
    "        np.copyto(roi_img, prev_img)\n",
    "        if cv.waitKey(0) == 13:\n",
    "            cv.imwrite(video[:-4] + \"-\" + str(no) + \"_roi.png\", roi_img)\n",
    "            return [x, y, a, a, rot]\n",
    "            break\n",
    "        elif cv.waitKey(0) == 113:  # shrink the roi\n",
    "            a = a-2\n",
    "        elif cv.waitKey(0) == 101:  # dilate the roi\n",
    "            a = a+2\n",
    "        elif cv.waitKey(0) == 95:  # move left\n",
    "            x = x-2\n",
    "        elif cv.waitKey(0) == 100:  # move right\n",
    "            x = x+2\n",
    "        elif cv.waitKey(0) == 115:  # move down\n",
    "            y = y+2\n",
    "        elif cv.waitKey(0) == 119:  # move up\n",
    "            y = y-2\n",
    "        elif cv.waitKey(0) == 120 or cv.waitKey(0) == 122:  # rotation\n",
    "            print(\"rotate!\")\n",
    "            if cv.waitKey(0) == 120:\n",
    "                rot += 0.2\n",
    "            elif cv.waitKey(0) == 122:\n",
    "                rot -= 0.2\n",
    "            M = cv.getRotationMatrix2D(rot_center, rot, 1)\n",
    "            roi_img = cv.warpAffine(roi_img, M, (WIDTH, HEIGHT))\n",
    "        # Draw the final roi square\n",
    "        cv.rectangle(roi_img, (x, y), (x + a, y + a), (0, 255, 0), 1)\n",
    "        cv.imshow(\"setROI2\", roi_img)\n",
    "\n",
    "\n",
    "def heat_pooling(img, N):\n",
    "    H, W = img.shape[:2]\n",
    "    heat_grid_img = np.zeros((H, W), dtype=np.uint8)\n",
    "    heat_grid = np.zeros((N, N), dtype=np.uint8)\n",
    "    G = int(H//N)\n",
    "    for y in range(N):\n",
    "        for x in range(N):\n",
    "            pos_sum = np.sum(img[G*y:G*(y+1), G*x:G*(x+1)])\n",
    "            heat_grid[y, x] = pos_sum\n",
    "            heat_grid_img[G*y:G*(y+1), G*x:G*(x+1)] = pos_sum\n",
    "    heat_grid_img = np.rot90(np.fliplr(heat_grid_img))\n",
    "    heat_grid_img = heat_grid_img/np.max(heat_grid_img)*255\n",
    "    heat_grid_img = heat_grid_img.astype(\"uint8\")\n",
    "\n",
    "    return heat_grid_img, heat_grid\n",
    "\n",
    "\n",
    "def heat_pooling_count(img, N=4):\n",
    "    H, W = img.shape[:2]\n",
    "    out = np.zeros((4, 4), dtype=np.uint8)\n",
    "    G = int(H//N)\n",
    "    for y in range(N):\n",
    "        for x in range(N):\n",
    "            out[y:y+1, x:x+1] = np.sum(img[G*y:G*(y+1), G*x:G*(x+1)])\n",
    "    return out[1:2, 1:2]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# main loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-11-24T04:40:43.817538Z",
     "start_time": "2022-11-24T04:40:43.775627Z"
    },
    "code_folding": []
   },
   "outputs": [],
   "source": [
    "def traceMouse(wkdir, videoname, no, f, f_end=0, croprot=[], inscale=0, roi=[]):\n",
    "    # Write log file\n",
    "    video = wkdir + videoname\n",
    "    now = datetime.now().strftime(\"%Y%m%d_%H%M%S\")\n",
    "    report = open(video[:-4] + \"-\" + str(no) + \".txt\", \"a\")\n",
    "    pos_report = open(video[:-4] + \"-\" + str(no) + \"_positions.txt\", \"a\")\n",
    "    report.write(\"Analyzed at: \"+now+\"\\n\")\n",
    "    report.write(\"Video name: \"+video+\"\\n\")\n",
    "    report.write(\"---------------------------------\\n\")\n",
    "    report.write(\"Global Parameters \\n\")\n",
    "    report.write(\"Frame frequency: \"+str(FF)+\"\\n\")\n",
    "    report.write(\"Binary threthold: \"+str(BIN_THR)+\"\\n\")\n",
    "    report.write(\"Vmin-max: \"+str(V_MIN)+str(V_MAX)+\"\\n\")\n",
    "    report.write(\"Duration: \"+str(DURATION)+\" frames\\n\")\n",
    "    report.write(\"---------------------------------\\n\")\n",
    "    capture = cv.VideoCapture(video)\n",
    "    vw = int(capture.get(cv.CAP_PROP_FRAME_WIDTH))\n",
    "    vh = int(capture.get(cv.CAP_PROP_FRAME_HEIGHT))\n",
    "    print(\"w: \" + str(vw) + \"h:\" + str(vh))\n",
    "\n",
    "    if roi == []:\n",
    "        x, y, w, h, rot = get_roi(capture, video, no, f)\n",
    "        scale = get_distance(ROIS[0], ROIS[1]) / 500\n",
    "    else:\n",
    "        x, y, w, h, rot, scale = roi\n",
    "\n",
    "    cv.destroyAllWindows()\n",
    "\n",
    "    # Report parameters for crop and rotation\n",
    "    print(\"Crop and Rotation: \", x, y, w, h, rot, scale)\n",
    "    report.writelines(\"crop and rotation: x=\"\n",
    "                      + str(int(x)) + \", y=\"\n",
    "                      + str(int(y)) + \", w=\"\n",
    "                      + str(int(w)) + \", h= \"\n",
    "                      + str(int(h)) + \", rot:\"\n",
    "                      + str(int(rot))\n",
    "                      + \"\\n\")\n",
    "\n",
    "    ct_ret1 = (int(h*0.25), int(w*0.25))\n",
    "    ct_ret2 = (int(h*0.75), int(w*0.75))\n",
    "\n",
    "    report.write(\"scale = \" + str(scale) + \" mm/pixel \\n\")\n",
    "    fcount = capture.get(cv.CAP_PROP_FRAME_COUNT)\n",
    "    track_img = np.ones((w, h, 3), dtype=np.uint8) * 255\n",
    "    for i in [int(0.25*w), int(0.5*w), int(0.75*w)]:\n",
    "        cv.line(track_img, (i, 1), (i, h), (192, 192, 192), 1)\n",
    "    for j in [int(0.25*h), int(0.5*h), int(0.75*h)]:\n",
    "        cv.line(track_img, (1, j), (w, j), (192, 192, 192), 1)\n",
    "\n",
    "    cv.rectangle(track_img, ct_ret1, ct_ret2, (128, 128, 128), 2)\n",
    "\n",
    "    heat_map = np.zeros((w, h), dtype=np.uint8)\n",
    "    cnts = []\n",
    "    centers = []\n",
    "    distance = 0\n",
    "    velmax = 0\n",
    "    stay = 0\n",
    "    v = 0\n",
    "    entry_cent_count = 0\n",
    "    beginPoint = (0, 0)\n",
    "    f_start = f\n",
    "\n",
    "    if f_end == 0:\n",
    "        f_end = f_start + DURATION\n",
    "\n",
    "    if capture.isOpened():\n",
    "        print(\"capture opened\")\n",
    "        while 1:\n",
    "            # End loop if enter is pressed\n",
    "            if cv.waitKey(1) == 13:\n",
    "                print(\"Enter is pressed, quit analysis.\")\n",
    "                cv.destroyAllWindows()\n",
    "                break\n",
    "\n",
    "            # End loop\n",
    "            if f >= f_end:\n",
    "                print(\"Finish and quit.\")\n",
    "                cv.destroyAllWindows()\n",
    "                break\n",
    "\n",
    "            # Process the image every FF drames\n",
    "            else:\n",
    "                rate = (f - f_start) / DURATION\n",
    "                capture.set(cv.CAP_PROP_POS_FRAMES, f)\n",
    "                ret, img = capture.read()\n",
    "\n",
    "                # RGB to grayscale\n",
    "                img = cv.cvtColor(img, cv.COLOR_RGB2GRAY)\n",
    "                img = cv.resize(img, (0, 0), fx=FX, fy=FY)\n",
    "                # Apply rotation and crop to image\n",
    "                (HEIGHT, WIDTH) = img.shape[:2]\n",
    "                rot_center = (WIDTH // 2, HEIGHT // 2)\n",
    "                M_img = cv.getRotationMatrix2D(rot_center, rot, 1)\n",
    "                rot_img = cv.warpAffine(img, M_img, (WIDTH, HEIGHT))\n",
    "                crop_img = rot_img[y:y + h, x:x + w]\n",
    "\n",
    "                # Get binary image\n",
    "                ret, bin_img = cv.threshold(\n",
    "                    crop_img, BIN_THR, 255, cv.THRESH_BINARY_INV)\n",
    "\n",
    "                # Remove small clusters which is not target\n",
    "                bin_img = cv.erode(bin_img,\n",
    "                                   np.ones((10, 10), np.uint8),\n",
    "                                   iterations=1)\n",
    "\n",
    "                cv.imshow(\"raw\", crop_img)\n",
    "                cv.moveWindow(\"raw\", 0, 0)\n",
    "                maxarea = 0\n",
    "                maxcnt = []\n",
    "                # Get the contours\n",
    "                contours = cv.findContours(bin_img, cv.RETR_EXTERNAL,\n",
    "                                           cv.CHAIN_APPROX_NONE)\n",
    "\n",
    "                if contours[0]:\n",
    "                    if len(contours[0]) == 1:\n",
    "                        maxcnt = contours[0][0]\n",
    "\n",
    "                    else:\n",
    "                        for i in range(0, len(contours[0])):\n",
    "                            cnt = contours[0][i]\n",
    "                            cnta = cv.contourArea(cnt)\n",
    "                            if cnta > maxarea:\n",
    "                                maxarea = cnta\n",
    "                                maxcnt = cnt\n",
    "                    try:\n",
    "                        mm = cv.moments(maxcnt)\n",
    "                        cx = mm[\"m10\"] / mm[\"m00\"]\n",
    "                        cy = mm[\"m01\"] / mm[\"m00\"]\n",
    "                    except:\n",
    "                        break\n",
    "\n",
    "                    centers.append([cx, cy])\n",
    "                    endPoint = (np.int(cx), np.int(cy))\n",
    "\n",
    "                    cv.drawContours(img, maxcnt, -1, (0, 0, 255), 5)\n",
    "                    cv.rectangle(crop_img, ct_ret1, ct_ret2, (0, 0, 0), 1)\n",
    "\n",
    "                    # Draw heatmap\n",
    "                    heat_map[endPoint] += 1\n",
    "                    if f % 100 == 0:\n",
    "\n",
    "                        cheat_map = heat_pooling(heat_map, HEAT_SIZE)[0]\n",
    "                        cheat_map = cv.applyColorMap(cheat_map, 15)\n",
    "                        cv.rectangle(cheat_map, ct_ret1,\n",
    "                                     ct_ret2, (128, 128, 128), 1)\n",
    "                        cv.imshow(\"Heat_map\", cheat_map)\n",
    "                        cv.moveWindow(\"Heat_map\", 3*w, 0)\n",
    "\n",
    "                    # Central area entry count\n",
    "                    if 0.25*w < endPoint[0] < 0.75*w and 0.25*w < endPoint[1] < 0.75*w:\n",
    "                        if 0.25*w < beginPoint[0] < 0.75*w and 0.25*w < beginPoint[1] < 0.75*w:\n",
    "                            pass\n",
    "                        else:\n",
    "                            entry_cent_count += 1\n",
    "\n",
    "                    # Instantaneous velocity calculation\n",
    "                    # Static state count\n",
    "                    if beginPoint != (0, 0):\n",
    "\n",
    "                        v = get_distance(beginPoint, endPoint)\n",
    "                        if V_MIN < v <= V_MAX:\n",
    "                            cv.line(track_img, beginPoint, endPoint,\n",
    "                                    (0 + int(rate * 255), 0,\n",
    "                                     255 - int(rate * 255)), 2)\n",
    "\n",
    "                            cv.line(bin_img, (np.int(cx), 0),\n",
    "                                    (np.int(cx), vh), (255, 255, 255), 1)\n",
    "                            cv.line(bin_img, (0, np.int(cy)),\n",
    "                                    (vw, np.int(cy)), (255, 255, 255), 1)\n",
    "                            distance += v\n",
    "                            if velmax < v:\n",
    "                                velmax = v\n",
    "\n",
    "                        else:\n",
    "                            stay += FF\n",
    "\n",
    "                        pos_report.writelines(\n",
    "                            str(f) + \": x=\" + str(int(cx)) + \",y=\" +\n",
    "                            str(int(cy)) + \",v=\" + str(int(v)) + \"\\n\")\n",
    "                    beginPoint = endPoint\n",
    "\n",
    "                    cv.putText(bin_img,\n",
    "                               \"Frame:\" +\n",
    "                               str(f) + \"-\"+str(int(rate * 100)) +\n",
    "                               \"%\", (10, 20),\n",
    "                               cv.FONT_HERSHEY_SIMPLEX, 0.6,\n",
    "                               (255, 255, 255), 1)\n",
    "                    cv.putText(\n",
    "                        bin_img,\n",
    "                        \"x=\" + str(int(cx)) + \",y=\" +\n",
    "                        str(int(cy)) + \",v=\" + str(int(v)), (10, 40),\n",
    "                        cv.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 1)\n",
    "\n",
    "                    cv.putText(bin_img, \"d=\" + str(int(distance))\n",
    "                               + \"e=\" + str(entry_cent_count)\n",
    "                               + \"s=\" + str(stay), (10, 60),\n",
    "                               cv.FONT_HERSHEY_SIMPLEX, 0.6,\n",
    "                               (255, 255, 255), 1)\n",
    "                    if f % 100 == 0:\n",
    "                        cv.imshow(\"Trajectory\", track_img)\n",
    "                        cv.moveWindow(\"Trajectory\", w, 0)\n",
    "\n",
    "                        cv.imshow(\"Binary image\", bin_img)\n",
    "                        cv.moveWindow(\"Binary image\", 2*w, 0)\n",
    "\n",
    "            f = f + FF\n",
    "\n",
    "        center_map4 = heat_pooling_count(heat_map)\n",
    "        center_time = np.sum(center_map4)\n",
    "        cv.imwrite(video[:-4] + \"-\" + str(no) + \"_trace.png\", track_img)\n",
    "        cv.imwrite(video[:-4] + \"-\" + str(no) + \"_heat.png\", cheat_map)\n",
    "        report.writelines(\n",
    "            \"distance(pixel) =\" + str(int(distance)) + \" pixel \\n\"\n",
    "            + \"distance(mm) =\" + str(int(distance) * scale) + \" mm \\n\"\n",
    "            + \"max velocity =\" + str(int(velmax) * scale * FF) + \" mm/s \\n\"\n",
    "            + \"stay =\" + str(int(stay)/V_FF) + \" s\\n\"\n",
    "            + \"stay to total time =\" + str(stay/DURATION*100) + \" %\\n\"\n",
    "            + \"central region time =\" +\n",
    "            str(int(center_time*FF)/V_FF) + \" s\\n\"\n",
    "            + \"central region time to total time = \" +\n",
    "            str(center_time*FF/DURATION*100)+\" %\\n\"\n",
    "            + \"entry central area =\" + str(entry_cent_count) + \" times\\n\"\n",
    "            + \"total recording time =\" + str(DURATION) + \" frames\\n\"\n",
    "        )\n",
    "\n",
    "        report.close()\n",
    "        pos_report.close()\n",
    "        print(\"center time = \"+str(center_time*FF)+\"frame\")\n",
    "        print(\"center time % = \"+str(center_time*FF/DURATION)+\"%\")\n",
    "        np.savetxt(video[:-4] + \"-\" + str(no) + \"_heat.txt\",\n",
    "                   heat_pooling(heat_map, HEAT_SIZE)[1], fmt='%d', delimiter=',')\n",
    "        centersvideo = wkdir + videoname\n",
    "        return centersvideo\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# analysis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "start_time": "2022-10-28T10:46:58.561Z"
    },
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "cts = traceMouse(\"D:/monkey/video/\", \"Video.mp4\", \"No\",\n",
    "                 1000, 0, 1)\n",
    "print(\"1 ok \")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.3 (default, Jul  2 2020, 17:30:36) [MSC v.1916 64 bit (AMD64)]"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {
    "height": "calc(100% - 180px)",
    "left": "10px",
    "top": "150px",
    "width": "383.993px"
   },
   "toc_section_display": true,
   "toc_window_display": true
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  },
  "vscode": {
   "interpreter": {
    "hash": "ca24d5afb66fcbb462fa71d0077b8bccae5a3d821a688de69b673156ae0e4c2b"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
